from openai import OpenAI
import os
import logging
from typing import List, Union

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class TextEmbedding:
    def __init__(self):
        """初始化OpenAI客户端，配置通义千问兼容模式"""
        self.api_key = os.getenv("DASHSCOPE_API_KEY")
        if not self.api_key:
            raise ValueError("请设置DASHSCOPE_API_KEY环境变量")

        # 初始化OpenAI客户端，指向通义兼容接口
        self.client = OpenAI(
            api_key=self.api_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )

        self.model = "text-embedding-v4"  # 通义text-embedding-v4模型
        self.dimensions = 1024  # 指定输出维度
        self.batch_size = 10  # 通义模型批处理大小限制

    def generate_embeddings(self, texts: Union[str, List[str]]) -> List[List[float]]:
        """
        生成文本嵌入向量
        :param texts: 单个文本字符串或文本列表
        :return: 嵌入向量列表
        """
        # 处理单文本输入
        if isinstance(texts, str):
            texts = [texts]

        # 过滤空文本并记录原始索引
        non_empty_texts = []
        original_indices = []
        for idx, text in enumerate(texts):
            cleaned_text = text.strip() if text else ""
            if cleaned_text:
                non_empty_texts.append(cleaned_text)
                original_indices.append(idx)

        # 存储结果的列表，保持原始输入顺序
        embeddings = [[] for _ in range(len(texts))]

        # 为空文本填充零向量
        for idx in range(len(texts)):
            if idx not in original_indices:
                embeddings[idx] = [0.0] * self.dimensions
                logger.warning(f"文本索引 {idx} 为空，返回零向量")

        # 如果没有有效文本，直接返回
        if not non_empty_texts:
            return embeddings

        # 分批处理文本
        try:
            for i in range(0, len(non_empty_texts), self.batch_size):
                batch = non_empty_texts[i:i + self.batch_size]
                batch_indices = original_indices[i:i + self.batch_size]

                # 调用嵌入API
                response = self.client.embeddings.create(
                    model=self.model,
                    input=batch,
                    dimensions=self.dimensions,
                    encoding_format="float"
                )

                # 验证返回结果数量
                if len(response.data) != len(batch):
                    raise ValueError(f"批量处理异常: 输入{len(batch)}个文本，返回{len(response.data)}个结果")

                # 提取嵌入向量并按原始索引放置
                for data, idx in zip(response.data, batch_indices):
                    if len(data.embedding) != self.dimensions:
                        raise ValueError(f"向量维度不匹配: 预期{self.dimensions}维，实际{len(data.embedding)}维")
                    embeddings[idx] = data.embedding

                logger.info(f"已处理 {min(i + self.batch_size, len(non_empty_texts))}/{len(non_empty_texts)} 个文本")

            return embeddings

        except Exception as e:
            logger.error(f"生成嵌入向量失败: {str(e)}")
            raise


# 示例用法
if __name__ == "__main__":
    try:
        embedder = TextEmbedding()

        # 测试单文本
        single_text = "如何在Boss直聘上优化企业招聘信息？"
        single_embedding = embedder.generate_embeddings(single_text)
        logger.info(f"单文本嵌入 - 维度: {len(single_embedding[0])}")
        logger.info(f"单文本嵌入 : {single_embedding[0]}")

        # 测试多文本（包含空文本）
        texts = [
            "求职者如何提高简历的曝光率？",
            "企业如何筛选合适的候选人？",
            "",  # 空文本测试
            "面试邀约的最佳时间是什么时候？",
            "如何撰写吸引人的职位描述？"
        ]
        batch_embeddings = embedder.generate_embeddings(texts)
        logger.info(f"批量嵌入 - 总数量: {len(batch_embeddings)}")
        logger.info(f"各向量维度: {[len(emb) for emb in batch_embeddings]}")

    except Exception as e:
        logger.error(f"测试失败: {str(e)}")
